import gym
import numpy as np
import pygame
import math

# 初始化Pygame
import torch
from gym.vector.utils import spaces
from torch import nn

# 游戏参数
screen_width = 400
screen_height = 400
cart_width = 60
cart_height = 20
pole_length = 100
pole_thickness = 10
cart_color = (0, 0, 0)
pole_color = (255, 0, 0)
background_color = (255, 255, 255)


class DQN(nn.Module):
    def __init__(self, input_size, output_size):
        super(DQN, self).__init__()
        self.fc1 = nn.Linear(input_size, 128)
        self.fc2 = nn.Linear(128, 64)
        self.fc3 = nn.Linear(64, output_size)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x


class InvertedPendulumEnv(gym.Env):
    def __init__(self):
        super(InvertedPendulumEnv, self).__init__()

        self.screen_width = 400
        self.cart_width = 60
        self.cart_range = 2.4  # 小车移动范围
        self.pole_length = 100
        self.pole_thickness = 10
        self.cart_speed = 0.1
        self.pole_angle = 0.1
        self.pole_angular_velocity = 0
        self.max_pole_angle = 15 * math.pi / 180  # 最大倾斜角度限制

        self.action_space = spaces.Discrete(2)  # 两个动作：向左或向右
        self.observation_space = spaces.Box(low=-np.inf, high=np.inf, shape=(4,))

        self.cart_x = self.screen_width // 2

    def reset(self):
        # 重置游戏状态
        self.cart_x = self.screen_width // 2
        self.pole_angle = 0.1
        self.pole_angular_velocity = 0
        return [self.cart_x, self.pole_angle, self.cart_x, self.pole_angular_velocity]

    def step(self, action):
        # 执行动作并返回下一个状态、奖励和是否完成的标志
        if action == 0:  # 向左移动
            self.cart_x -= self.cart_speed
            self.pole_angle += 0.0005
        elif action == 1:  # 向右移动
            self.cart_x += self.cart_speed
            self.pole_angle -= 0.0005
        done = (self.pole_angle > self.max_pole_angle or self.pole_angle < -self.max_pole_angle or
                self.cart_x < self.screen_width // 2 - self.cart_range * self.cart_width or
                self.cart_x > self.screen_width // 2 + self.cart_range * self.cart_width)
        reward = 10 / (0.1+self.pole_angle) if not done else 0
        return [self.cart_x, self.pole_angle, reward], done

    def show(self):
        screen.fill(background_color)

        # 绘制小车
        pygame.draw.rect(screen, cart_color,
                         (self.cart_x - cart_width // 2, screen_height - cart_height, cart_width, cart_height))

        # 计算杆子的端点坐标
        pole_x = self.cart_x + pole_length * math.sin(self.pole_angle)
        pole_y = screen_height - cart_height - pole_length * math.cos(self.pole_angle)

        # 绘制杆子
        pygame.draw.line(screen, pole_color, (self.cart_x, screen_height - cart_height), (pole_x, pole_y),
                         pole_thickness)

        pygame.display.update()
    def run(self):
        self.pole_angle += self.pole_angular_velocity
        self.pole_angular_velocity = 0.000398 * math.sin(self.pole_angle)

pygame.init()

env = InvertedPendulumEnv()
# 初始化屏幕
screen = pygame.display.set_mode((screen_width, screen_height))
pygame.display.set_caption("Inverted Pendulum Game")

input_size = 4
output_size = 2
model = DQN(input_size, output_size)
model.load_state_dict(torch.load('rlmodel(2).pth', map_location=torch.device('cpu')))
model.eval()  # 设置模型为评估模式
# 游戏参数
cart_speed = 0.1  # 小车速度

# 游戏循环
running = True
while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
    env.run()
    keys = pygame.key.get_pressed()
    if keys[pygame.K_LEFT]:
        state, done = env.step(0)
    if keys[pygame.K_RIGHT]:
        state, done = env.step(1)
    # 渲染游戏界面
    env.show()
# 退出游戏
pygame.quit()
